import cv2
import numpy as np
from collections import defaultdict
import math



def orthogonal_regression(points):
    # 计算点集的中心
    center = np.mean(points, axis=0)

    # 将点集移到中心
    centered_points = points - center

    # 计算协方差矩阵
    cov_matrix = np.dot(centered_points.T, centered_points) / (len(points) - 1)

    # 使用特征值分解求解协方差矩阵的特征向量
    eigenvalues, eigenvectors = np.linalg.eig(cov_matrix)

    # 找到最小特征值对应的特征向量
    min_eigenvalue_index = np.argmin(eigenvalues)
    min_eigenvector = eigenvectors[:, min_eigenvalue_index]

    # 通过特征向量计算回归线的斜率和截距
    slope = min_eigenvector[1] / min_eigenvector[0]
    intercept = center[1] - slope * center[0]

    return slope, intercept


def least_squares_fit(points):
    # 提取点的 x 和 y 坐标
    x = points[:, 1]
    y = points[:, 0]

    # 使用最小二乘法计算直线的斜率和截距
    A = np.vstack([x, np.ones(len(x))]).T
    m, q = np.linalg.lstsq(A, y, rcond=None)[0]

    return m, q


def Least_squares(points):
    x=points[:,1]
    y=points[:,0]
    x_ = x.mean()
    y_ = y.mean()
    m = np.zeros(1)
    n = np.zeros(1)
    k = np.zeros(1)
    p = np.zeros(1)
    for i in np.arange(len(x)):
        k = (x[i]-x_) * (y[i]-y_)
        m += k
        p = np.square(x[i]-x_)
        n = n + p
    a = m/n
    b = y_ - a * x_
    return a, b


def only_one_line(points):
    # 合并成一条线
    merged_dict = defaultdict(list)

    for first_dim, second_dim in points:
        merged_dict[first_dim].append(second_dim)

    result = np.array([[key, sum(values) / len(values)]
                      for key, values in merged_dict.items() if len(values)==4])

    return result

def calculate_slope(image, is_return_process_image=False):
    """
    计算斜率。

    参数：
    - image：输入图像，必须是BGR格式的彩色图像。
    - is_return_process_image：是否返回处理后的图像，默认为False。

    返回值：
    - 如果is_return_process_image为True，则返回一个元组，包含以下内容：
        - done：布尔值，表示计算斜率是否成功。
        - slope：斜率值。
        - binary_cleaned：处理后的二值图像。
    - 如果is_return_process_image为False，则返回一个元组，包含以下内容：
        - done：布尔值，表示计算斜率是否成功。
        - slope：斜率值。

    注意：
    - 输入图像会进行阈值处理、掩码操作、开闭运算等处理步骤。
    - 最终会使用Canny边缘检测来获取图像中的边缘点。
    - 边缘点会经过正交回归或最小二乘法拟合来计算斜率和截距值。

    示例用法：
    ```
    image = cv2.imread('image.jpg')
    done, slope = caculate_slope(image)
    ```
    """

    hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)

    # 创建掩膜
    lower = np.array([0, 0, 0])
    upper = np.array([179, 255, 49])
    mask = cv2.inRange(hsv, lower, upper)

    # cv2.imshow('result', result)
    # cv2.waitKey(0)

    kernel_size = 7

    # 使用开运算去除小的白色噪点
    opening_kernel = cv2.getStructuringElement(
        cv2.MORPH_RECT, (kernel_size, kernel_size))
    binary_opened = cv2.morphologyEx(mask, cv2.MORPH_OPEN, opening_kernel)

    # 使用闭运算填充小的黑色孔洞
    closing_kernel = cv2.getStructuringElement(
        cv2.MORPH_RECT, (kernel_size, kernel_size))
    binary_cleaned = cv2.morphologyEx(
        binary_opened, cv2.MORPH_CLOSE, closing_kernel)
    
    # cv2.namedWindow("Binary Image", cv2.WINDOW_NORMAL)
    # cv2.resizeWindow("Binary Image", 680, 640)
    # cv2.imshow("Binary Image", binary_cleaned)
    # cv2.waitKey(0)
    # cv2.imwrite('binary_image.jpg', binary_cleaned)
    
    # 手动指定要更改的像素位置范围
    # 例如，设定 x 坐标从 100 到 200，y 坐标从 50 到 150 的区域为白色
    # start_x, end_x = 300, 380
    # start_y, end_y = 440, 520

    # # 将指定区域设置为白色（255）
    # binary_cleaned[start_y:end_y, start_x:end_x] = 0

    # cv2.imshow("Binary Image", binary_cleaned)
    # cv2.waitKey(0)

    # 应用 Canny 边缘检测
    binary_cleaned = cv2.Canny(binary_cleaned, 150, 450)  # 可以调整参数来适应不同的图像
    # cv2.imwrite('binary_image.jpg', binary_cleaned)

    points = np.argwhere(binary_cleaned != 0)
    points = only_one_line(points)

    

    
    if isinstance(points, np.ndarray) and points.ndim == 2 and len(points) > 1:
        slope, intercept = Least_squares(points)
        # slope, intercept = orthogonal_regression(points)
        done=True
    else:
        slope=0
        intercept=0
        done=False



    if is_return_process_image:
        # 绘制拟合直线的端点
        if not np.isnan(slope):
            x1 = 0
            y1 = int(slope * x1 + intercept)
            x2 = binary_cleaned.shape[1]
            y2 = int(slope * x2 + intercept)

            # 绘制直线
            cv2.line(binary_cleaned, (x1, y1), (x2, y2), (255, 255, 255), 2)

            result_rad = math.atan(slope)
            result_deg = math.degrees(result_rad)
            slope = result_deg
        else:
            slope = 90
        return done, slope, binary_cleaned
    else:
        if not np.isnan(slope):
            result_rad = math.atan(slope)
            result_deg = math.degrees(result_rad)
            slope = result_deg
        else:
            slope = 90
        return done, slope

if __name__=='__main__':

    # image=cv2.imread('gym_pih\envs\experiment6.jpg')
    import platform
    if platform.system() == 'Windows':
        from gym_pih.hardware.Imagingsource_for_windows import Camera as Camera
    elif platform.system() == 'Linux':
        from gym_pih.hardware.Imagingsource_for_linux import Camera as Camera
    from gym_pih.hardware.LightSource import LightSource
    light = LightSource()
    light.Light_Open(6)
    camera= Camera("1280x960")
    camera.Open()
    if platform.system() == 'Windows':
        camera.SetProperty("Gain","Value",5.2)
        camera.SetProperty("Exposure","Value",float(1/60))
    while cv2.waitKey(1) & 0xFF != 27:
        _,image=camera.GetImage()
        _,slope,binary_cleaned=calculate_slope(image,is_return_process_image=True)
        # print(binary_cleaned.shape)

        print(slope, "°")
        # 显示图像
        cv2.namedWindow("Image", cv2.WINDOW_NORMAL)
        # cv2.resizeWindow("Image", 770, 590)
        cv2.imshow("Image", image)
        cv2.namedWindow("Fitted Line", cv2.WINDOW_NORMAL)
        # cv2.resizeWindow("Fitted Line", 770, 590)
        cv2.imshow("Fitted Line", binary_cleaned)
        # cv2.imwrite('binary_image.jpg', binary_cleaned)
        # cv2.imwrite('image.jpg', image[130:,250:1020])

    
    cv2.destroyAllWindows()
    camera.Close()
    light.Light_Close()
    light.port_close()
